import pandas as pd
import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.dates as mdates
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter
import cartopy.feature as cfeature
import shapely.geometry as sgeom
from datetime import datetime
import glob
from tqdm import tqdm
# 自定义label
from matplotlib.patches import Patch
from matplotlib.lines import Line2D
from gsj_typhoon import split_str_id,tydat_CMA,tydat
import cartopy.io.shapereader as shpreader
import geopandas as gpd
from typlot.config.global_config import *

names = ['mojie_28','dusurui_16','gaemi_09','haikui_38','kangni_54','shantuo_44','saola_25','koinu_49']
# names = ['dusurui_16']
tynames,tyids = split_str_id(names)
draw_opt = False


for tyname,tyid in zip(tynames,tyids) :
    RIstd = 7
    path = os.path.join(global_ensdir, f"{tyname}_{tyid}")
    dates= sorted(os.listdir(path))
    paths=[os.path.join(path,n) for n in dates]
    pathobs = os.path.join(global_obsdir, f"{tyname}_CMAobs.txt")
    pic_savepath = os.path.join( global_picdir ,"w_p_track",f"{tyname}")
    os.makedirs(pic_savepath,exist_ok=True)
    
    def basemap(ax, title=None, extend_degrees=15):
        """
        只显示中国省界的底图
        
        推荐使用: admin_1_states_provinces (面状数据)
        """
        
        ax.add_feature(cfeature.COASTLINE, linewidth=0.5)
        ax.add_feature(cfeature.BORDERS, linestyle=':', linewidth=0.5)
        
        # ✅ 使用 admin_1_states_provinces (面状数据，更完整)
        try:
            provinces_shp = shpreader.natural_earth(
                resolution='10m',
                category='cultural',
                name='admin_1_states_provinces'  # 面状，不是 _lines
            )
            
            gdf = gpd.read_file(provinces_shp)
            
            # 筛选中国省份
            china_provinces = gdf[
                (gdf['admin'] == 'China') | 
                (gdf['adm0_a3'] == 'CHN')
            ]
            
            # 绘制省界
            ax.add_geometries(
                china_provinces.geometry,
                ccrs.PlateCarree(),
                facecolor='none',
                edgecolor='black',
                linewidth=0.5
            )
            
            print(f"✅ 成功绘制 {len(china_provinces)} 个中国省份")
            
        except Exception as e:
            print(f"❌ 加载失败: {e}")
        
        # 网格线
        gl = ax.gridlines(
            draw_labels=True, 
            alpha=0.3, 
            linestyle='--', 
            linewidth=0.5,
            color='gray'
        )
        
        gl.top_labels = False
        gl.right_labels = False
        gl.xlabel_style = {'size': 8}
        gl.ylabel_style = {'size': 8}
        gl.xformatter = LongitudeFormatter(number_format='.1f')
        gl.yformatter = LatitudeFormatter(number_format='.1f')
        
        if title:
            ax.set_title(title)
    
    
    def track(ax, ty, sep=1, color='Black', alpha=0.7, linewidth=1.5):
        lon = np.array(ty.lon, dtype=float)
        lat = np.array(ty.lat, dtype=float)
        
        if len(lon) <= 1:
            return
        
        # 经度连续化处理
        lon_continuous = lon.copy()
        
        for i in range(1, len(lon_continuous)):
            diff = lon_continuous[i] - lon_continuous[i-1]
            
            if diff > 180:
                lon_continuous[i:] -= 360
            elif diff < -180:
                lon_continuous[i:] += 360
        
        # 🔧 统一坐标系转换
        # 选择一个统一的坐标系，与 ax 的投影保持一致
        data_crs = ccrs.PlateCarree()  # 数据坐标系
        
        # 绘制轨迹点
        ax.scatter(lon, lat, c=color, s=15, transform=data_crs, 
                   alpha=alpha, zorder=5, edgecolors='white', linewidths=0.3)
        
        # 🔧 绘制连续的轨迹线（使用相同的坐标系）
        ax.plot(lon_continuous, lat, color=color, linewidth=linewidth, 
                alpha=alpha, transform=data_crs, zorder=4)
        
        # 图例
        if not hasattr(ax, '_legend_created'):
            legend_elements = [
                Line2D([0], [0], color='black', lw=1, alpha=0.5, label='ensemble members', marker='o'),
                Line2D([0], [0], color='lime', lw=2, label='mean of members'),
                Line2D([0], [0], color='red', lw=2, label='Analysis'),
                Line2D([0], [0], color='blue', lw=2, label='observation'),
            ]
            ax.legend(handles=legend_elements, loc='lower right', fontsize=8)
            ax._legend_created = True
    
    
    def SLP(ax,ty,title,color="Black"):
        ax.plot(ty.time, ty.pmin, 'o-', color=color)
        # legend
        legend_elements = [
            Line2D([0], [0], color='black', lw=1, label='emsemble members',marker='o',markersize=5,ls="none"),
            Line2D([0], [0], color='green', lw=1, label='mean of members',marker='o',markersize=5,ls="none"),
            Line2D([0], [0], color='red', lw=1, label='analysis',marker='o',markersize=5,ls="none"),
            Line2D([0], [0], color='blue', lw=1, label='observation',marker='o',markersize=5,ls="none"),
            ]
        ax.legend( handles=legend_elements,loc='lower right',fontsize=10 )    
        # label
        min_slp_label = np.arange(920,1030,10)
        ax.set_title(title)
        ax.set_xlabel('Time (month-day)')
        ax.set_ylabel('min SLP (hPa)')
        ax.set_yticks( min_slp_label )
        ax.tick_params(axis='y')
        ax.xaxis.set_major_locator(mdates.DayLocator(interval=2))  # 每5天一个主刻度
        ax.xaxis.set_minor_locator(mdates.DayLocator(interval=1))  # 每1天一个次刻度
        ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d"))  # 主刻度显示的格式
    
    
    def maxw(ax,ty,title,color="Black"):
        ax.plot(ty.time, ty.umax, 'o-', color=color)
        # legend
        legend_elements = [
            Line2D([0], [0], color='black', lw=1, label='emsemble members',marker='o',markersize=5,ls="none"),
            Line2D([0], [0], color='green', lw=1, label='mean of members',marker='o',markersize=5,ls="none"),
            Line2D([0], [0], color='red', lw=1, label='analysis',marker='o',markersize=5,ls="none"),
            Line2D([0], [0], color='blue', lw=1, label='observation',marker='o',markersize=5,ls="none"),
            ]
        ax.legend( handles=legend_elements,loc='upper right',fontsize=10 )   
        # label
        max_wind_label = np.arange(0,70,10)
        ax.set_title(title)
        ax.set_xlabel('Time (month-day)')
        ax.set_ylabel('max wind (m/s)')
        ax.set_yticks( max_wind_label )
        ax.tick_params(axis='y')
        ax.xaxis.set_major_locator(mdates.DayLocator(interval=2))  # 每5天一个主刻度
        ax.xaxis.set_minor_locator(mdates.DayLocator(interval=1))  # 每1天一个次刻度
        ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d"))  # 主刻度显示的格式
        
        
    
    

    
    
    for i in tqdm(range(len(paths)),desc = f'processing {tyname}'):
        path=paths[i]
        title=dates[i]+tyname
        # mean
        meanpath = glob.glob(os.path.join(path,f"*{tyid}_mean"))[0]
        # reanalysis
        n0path = os.path.join(path, "TRACK_ID_0")

        # other path
        start = 'TRACK_ID'   # 获取目录下所有文件和子目录
        all_files = os.listdir(path)   # 过滤出文件且符合后缀要求
        filtered_files = [f for f in all_files if f.startswith(start) and f != "TRACK_ID_0"]   # 打印所有符合后缀的文件名
        
    ##################  track #################
        fig,ax= plt.subplots(subplot_kw={'projection': ccrs.PlateCarree(central_longitude=150)},dpi=300,figsize=(10,5))
        #绘制track 图
        for f in filtered_files:
            ty=tydat(os.path.join(path,f),RIstd)
            track(ax,ty,color="Black")
        # 绘制mean
        ty=tydat(meanpath,RIstd,skiprows=5)  
        track(ax,ty,color="lime")
        # 绘制trackid 0
        ty=tydat(n0path,RIstd,skiprows=2)  
        track(ax,ty,color="Red")
        # 绘制实况
        tyobs = tydat_CMA(pathobs)
        track(ax,tyobs,color="blue")
        basemap(ax,title=title)  
        if draw_opt == False:
            plt.savefig(os.path.join(pic_savepath,f"track_{title}"))
            plt.close()
        else : 
            plt.show()
        
    ################## SLP  #################    
        fig,ax= plt.subplots(dpi=300,figsize=(10,5))
        #绘制集合成员
        
        for f in filtered_files:
            ty = tydat(os.path.join(path,f),RIstd)
            SLP(ax,ty,title,color="Black")
        # 绘制mean
        ty=tydat(meanpath,RIstd,skiprows=5)  
        SLP(ax,ty,title,color="lime") 
        # 绘制trackid 0
        ty=tydat(n0path,RIstd,skiprows=2)  
        SLP(ax,ty,title,color="Red")
        # 绘制 实况
        SLP(ax,tyobs,title,color="blue")
        if draw_opt == False:
            plt.savefig( os.path.join(pic_savepath,f'SLP_{title}') )
            plt.close()
        else : 
            plt.show()
    
    #################  maxw   #################
        fig,ax= plt.subplots(dpi=300,figsize=(10,5));
        #绘制集合成员
        for f in filtered_files:
            ty = tydat(os.path.join(path,f),RIstd)
            maxw(ax,ty,title,color="Black")
        # 绘制mean
        ty=tydat(meanpath,RIstd,skiprows=5)  
        maxw(ax,ty,title,color="lime") 
        # 绘制trackid 0
        ty=tydat(n0path,RIstd,skiprows=2)  
        maxw(ax,ty,title,color="Red")
        #  绘制实况
        maxw(ax,tyobs,title,color="blue")
        if draw_opt == False:
            plt.savefig( os.path.join(pic_savepath,f'maxw_{title}') )
            plt.close()
        else:
            plt.show()